import torch
import clip
from PIL import Image

# 加载CLIP模型
device = "cuda" if torch.cuda.is_available() else "cpu"
model, preprocess = clip.load("ViT-B/32", device=device)

def load_and_preprocess(image_path):
    image = Image.open(image_path)
    return preprocess(image).unsqueeze(0).to(device)

# 加载正负样本
positive_img = load_and_preprocess("positive.jpg")
negative_img = load_and_preprocess("negative.jpg")

# 提取特征并归一化
with torch.no_grad():
    positive_feature = model.encode_image(positive_img)
    negative_feature = model.encode_image(negative_img)
    
# 归一化特征
positive_feature /= positive_feature.norm(dim=-1, keepdim=True)
negative_feature /= negative_feature.norm(dim=-1, keepdim=True)

def classify_image(test_path):
    # 处理测试图片
    test_img = load_and_preprocess(test_path)
    with torch.no_grad():
        test_feature = model.encode_image(test_img)
    test_feature /= test_feature.norm(dim=-1, keepdim=True)
    
    # 计算余弦相似度
    sim_positive = (test_feature @ positive_feature.T).item()
    sim_negative = (test_feature @ negative_feature.T).item()
    
    return "存在目标零件" if sim_positive > sim_negative else "无目标零件"

# 示例使用
test_result = classify_image("test_image1.jpg")
print(test_result)

test_result = classify_image("test_image2.jpg")
print(test_result)

test_result = classify_image("test_image3.jpg")
print(test_result)

test_result = classify_image("test_image4.jpg")
print(test_result)
